import json
import os
import shutil

import cv2

from exchange_opt import *


# blue 亮 200 暗 中
# red 亮 175 暗 中

image_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/goo 分批收集/李浩淼/wrong"

# 获取文件夹A中所有图片文件路径
image_files = [os.path.join(image_path, f) for f in os.listdir(image_path) if
               f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

# 根据文件名进行排序
sorted_image_files = sorted(image_files, key=lambda x: os.path.basename(x))

# sorted_image_files = [r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/red/Video_20250113104416448/exchange_red_Video_20250113104416448_27.jpg"]


# 定义类别和对应的group_id
category_to_group_id = {
    "box": 0,
    "arrows_right": 1,
    "arrows_left": 2
}

# 打印排序后的图片文件路径
for file_path in sorted_image_files:
    print(file_path)

    rst = []

    classes = {'box': 0, 'arrow_right': 1, 'arrow_left': 2}

    image = cv2.imread(file_path)
    image_height, image_width, _ = image.shape

    # cv2.line(image, (0, 0), (image_width, 00), (0, 0, 0), 10)
    # cv2.line(image, (image_width, 0), (image_width, image_height), (0, 0, 0), 10)
    # cv2.line(image, (image_width, image_height), (0, image_height), (0, 0, 0), 10)
    # cv2.line(image, (0, image_height), (0, 0), (0, 0, 0), 10)

    json_file_path = file_path.replace(".jpg", ".json")

    with open(json_file_path, 'r') as file:
        data = json.load(file)

    for shape in data["shapes"]:
        class_id = category_to_group_id.get(shape["label"])  # shape["label"] 类别名

        if shape["shape_type"] == "rectangle":
            # to xywh
            # find left/right corner
            x1, y1, x2, y2 = 0, 0, 0, 0
            if shape["points"][0][0] < shape["points"][1][0]:
                x1 = shape["points"][0][0]
                x2 = shape["points"][1][0]
            else:
                x1 = shape["points"][1][0]
                x2 = shape["points"][0][0]

            if shape["points"][0][1] < shape["points"][1][1]:
                y1 = shape["points"][0][1]
                y2 = shape["points"][1][1]
            else:
                y1 = shape["points"][1][1]
                y2 = shape["points"][0][1]

            width = abs(shape["points"][1][0] - shape["points"][0][0])
            height = abs(shape["points"][1][1] - shape["points"][0][1])
            # print(width, height)
            width = abs(x1 - x2)
            height = abs(y1 - y2)
            # print(width, height)

            keypoints = []
            for keypoint_label in [f"{shape['label']}_{i}" for i in range(1, 7)]:
                keypoint = next((keypoint for keypoint in data["shapes"] if keypoint["label"] == keypoint_label), None)
                if keypoint:
                    keypoint_x = keypoint["points"][0][0]
                    keypoint_y = keypoint["points"][0][1]
                    keypoint_visibility = 1.0  # Assuming all keypoints are visible 0：没有明显露出，不可见 1：被遮挡，不可见 2：可见
                    keypoints.extend([keypoint_x, keypoint_y, keypoint_visibility])
                else:
                    keypoints.extend([0, 0, 0])  # If keypoint not found, use 0 values

            rst.append([
                x1, y1, width, height,
                class_id, 1.0,
                keypoints[0], keypoints[1], keypoints[2],
                keypoints[3], keypoints[4], keypoints[5],
                keypoints[6], keypoints[7], keypoints[8],
                keypoints[9], keypoints[10], keypoints[11],
                keypoints[12], keypoints[13], keypoints[14],
                keypoints[15], keypoints[16], keypoints[17]
            ])

    if len(rst) == 0:
        print("rst == 0")
        continue



    for rs in rst:
        class_id = rs[4]
        color = (0, 255, 255)
        if class_id == 0:
            color = (0, 0, 255)
        elif class_id == 1:
            color = (255, 0, 255)
        cv2.circle(image, (int(rs[6]), int(rs[7])), 5, color, 2)
        cv2.putText(image, "1", (int(rs[6]), int(rs[7])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.circle(image, (int(rs[9]), int(rs[10])), 5, color, 2)
        cv2.putText(image, "2", (int(rs[9]), int(rs[10])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.circle(image, (int(rs[12]), int(rs[13])), 5, color, 2)
        cv2.putText(image, "3", (int(rs[12]), int(rs[13])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.circle(image, (int(rs[15]), int(rs[16])), 5, color, 2)
        cv2.putText(image, "4", (int(rs[15]), int(rs[16])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.circle(image, (int(rs[18]), int(rs[19])), 5, color, 2)
        cv2.putText(image, "5", (int(rs[18]), int(rs[19])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.circle(image, (int(rs[21]), int(rs[22])), 5, color, 2)
        cv2.putText(image, "6", (int(rs[21]), int(rs[22])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)

        cv2.rectangle(image, (int(rs[0]), int(rs[1])), (int(rs[0] + rs[2]), int(rs[1] + rs[3])), color, 2)
        cv2.putText(image, str(class_id),  (int(rs[0]), int(rs[1])), cv2.FONT_HERSHEY_SIMPLEX, 2.0,color, 3)


    cv2.namedWindow("image", 0)
    cv2.imshow("image", image)

    key = cv2.waitKey(0)
    if key == 27:
        break
    elif key == ord('c'):
        print("press c")
        output_folder = os.path.join(os.path.dirname(file_path), "wrong")
        os.makedirs(output_folder, exist_ok=True)
        # 获取文件名
        file_name = os.path.basename(json_file_path)
        # 构造目标路径
        destination_file = os.path.join(output_folder, file_name)
        # 移动文件
        shutil.move(json_file_path, destination_file)
        # 获取文件名
        file_name = os.path.basename(file_path)
        # 构造目标路径
        destination_file = os.path.join(output_folder, file_name)
        # 移动文件
        shutil.move(file_path, destination_file)

cv2.destroyAllWindows()
